import collections
import json
import math
import time
from typing import List, Optional, Dict, Callable
from pprint import pprint

from gai.v2.spark.external.client_base import ClientBase, simplify_data, md5sum, generate_secret
from gai.v2.spark.external import retriable_execute
from gai.v2.spark.external.generic import equal_subranges


def get_table(day: str):
    """
    Generates the table name.

    Args:
        day:
            calendar date in ``YYYYMMDD`` format

    Returns:
        the name of the table to look up into.
    """
    return 'gz_ft:gz_ft_' + day[0:6]


def get_column_family(day):
    """
    Generates the column family.

    Args:
        day:
            calendar date in ``YYYYMMDD`` format

    Returns:
        the column family.
    """
    tmp = math.floor((int(day[6:8]) - 1) / 7)
    return {
        0: 'w0',
        1: 'w1',
        2: 'w2',
        3: 'w3',
        4: 'w4',
    }.get(tmp, 'error')


def get_column_qualifier(day):
    """
        Generates the column qualifier.

    Args:
        day:
            calendar date in ``YYYYMMDD`` format

    Returns:
        the column qualifier.

    """
    tmp = (int(day[6:8])) % 7
    return {
        1: 'd1',
        2: 'd2',
        3: 'd3',
        4: 'd4',
        5: 'd5',
        6: 'd6',
        0: 'd7',
    }.get(tmp, 'error')


def make_column_spec(day):
    """
        Generates the column specification.

    Args:
        day:
            calendar date in ``YYYYMMDD`` format

    Returns:
        the column spec.
    """
    return get_column_family(day) + ':' + get_column_qualifier(day)


def encode_gid(gid):
    """
    Converts the gid to row key.

    Args:
        gid:

    Returns:
        the row key.
    """
    if len(gid) > 4:
        return gid[- 4:] + gid[0:- 4]
    else:
        return gid


def decode_rowkey(rowkey):
    """
    Converts the row key to gid.

    Args:
        rowkey:

    Returns:
        the gid.
    """
    if len(rowkey) > 4:
        return rowkey[4:] + rowkey[:4]
    else:
        return rowkey


def make_body(table, gets, attributes):
    return {
        "table"     : table,
        "gets"      : gets,
        "attributes": attributes
    }


QueryPoint = collections.namedtuple('QueryPoint', ['day', 'gid'])


def make_body_from_points(points: List[QueryPoint], features: List[str], encode_gid: Callable,
                          extra_params=None) -> json:
    def make_get(point: QueryPoint):
        return {
            "row"    : encode_gid(point.gid),
            "columns": [make_column_spec(point.day)]
        }

    table = get_table(points[0].day)
    gets = [make_get(p) for p in points]
    attr_features = ",".join(features)
    attributes = {"features": attr_features}
    if extra_params is not None:
        attributes["extra_params"] = extra_params

    body = make_body(table, gets, attributes)
    return json.dumps(body)


def make_exception(points, features, response_str):
    return RuntimeError(
        "The query is unsuccessful. Here is the response body: {}. "
        "And here is the query: points = {}, features = {}".format(
            response_str, points, features))


class FeatureClient(ClientBase):
    emptystr_placeholder = "empty-string-placeholder"
    nullstr_placeholder = "null-string-placeholder"

    def __init__(self, url, user_name, secret, accept_emptystr=True):
        """

        Args:
            url:
                the url of the service.
                For example, <code>http://localhost:5081/gt-feature-data-service</code>
            user_name:
                the name of the client user.
            secret:
                the password assigned to the user.
            accept_emptystr:
                indicator whether an empty string is acceptable.
        """
        super(FeatureClient, self).__init__(url, user_name, secret)
        self._accept_emptystr = accept_emptystr

    def get(self, gid: str, day: str, features: List[str], extra_params=None) -> Optional[str]:
        """
        Given a gid, a date and a list of feature names, find the corresponding string representation
        of feature values.

        Args:
            gid:
            day:
                calendar date in ``YYYYMMDD`` format.
            features:
                a list of feature names

        Returns:
            a string representing the found features, or ``None``
        """
        batch_result = self.batch_get([gid], day, features, extra_params)
        column_spec = make_column_spec(day)
        return batch_result[0]["cells"].get(column_spec)

    def batch_get(self, gids: List[str], day: str, features: List[str], extra_params=None) -> list:
        """
        Given a list of gids, a date and a list of feature names, find the corresponding list of string presentation
        of feature values.

        Args:
            gids:
            day:
                calendar date in ``YYYYMMDD`` format.
            features:
                a list of feature names

        Returns:
            the list of matched features. Mismatched records are given ``None`` as placeholders.
        """

        points = [QueryPoint(day, gid) for gid in gids]
        body = make_body_from_points(points, features, self.encode_gid, extra_params)
        return self._execute_with_body(body,
                                       lambda x: make_exception(points, features, x))

    def batch_get_multi(self, points: List[QueryPoint],
                        features: List[str], max_retries=3, sleep_secs=1, extra_params=None) -> list:
        """
        Given a list of query points and a list of feature names, return the
        list of request features for the query points.

        Args:
            points:
                a list of query points.
            features:
                a list of feature names.
        Returns:
            the list of matched features. Mismatched records are given
            ``None`` as placeholders.
        """

        assert len(points) > 0 and max_retries >= 0

        def equal_table(x: QueryPoint, y: QueryPoint):
            return x.day[0:6] == y.day[0:6]

        result = []
        for subrange in equal_subranges(points, equal_table):
            assert len(subrange) > 0
            body = make_body_from_points(subrange, features, self.encode_gid, extra_params)
            task = lambda: self._execute_with_body(body,
                                                   lambda x: make_exception(subrange, features, x))
            batch_result = retriable_execute(task,
                                             max_tries=max_retries + 1,
                                             sleep_secs=sleep_secs)
            result.extend(batch_result)
        return result

    def _execute_with_body(self, body, make_exception: Callable):
        """Given the query body and an exception constructor,
        query with the body, and return the result.
        In case of exception, raise the exception constructed
        with ``make_exception``.

        Args:
            body:
                the query body.
            make_exception:
                a unary function that create an Exception object.

        Returns:
            the query result.
        """

        response = self.send(body)
        response_str = response.data.decode()
        response_json = json.loads(response_str)
        if response_json.get("result", -1) == 0:
            simp_data = simplify_data(response_json["data"])
            for i in range(len(simp_data)):
                simp_data[i]["row"] = self.decode_rowkey(simp_data[i]["row"])
            return simp_data
        else:
            raise make_exception(response_str)

    def get_headers(self, body) -> Dict[str, str]:
        t = time.time()
        timestamp = str(round(t * 1000))
        appid = self._user_name
        secret = self._secret
        serviceName = "gt-feature-data-service"
        sign = md5sum(
            md5sum(serviceName + appid) + md5sum(body) + secret + timestamp
        )

        headers = {
            'Timestamp'  : timestamp,
            "App-id"     : appid,
            'Sign'       : sign,
            'serviceName': serviceName,
            # 'userId'     : '57'  # for debug use
        }

        return headers

    def encode_gid(self, x):
        if x is None:
            return FeatureClient.nullstr_placeholder
        elif self.accept_emptystr() and len(x) == 0:
            return FeatureClient.emptystr_placeholder
        else:
            return encode_gid(x)

    def decode_rowkey(self, x):
        if x == FeatureClient.nullstr_placeholder:
            return None
        elif self.accept_emptystr() and x == FeatureClient.emptystr_placeholder:
            return ""
        else:
            return decode_rowkey(x)

    def accept_emptystr(self):
        return self._accept_emptystr


def demo_feature_client_troubleshooting(url, user_name, secret):
    def print_result(index, result):
        print("{0} {1} {0}".format("-" * 10, index))
        pprint(result)

    client = FeatureClient(url, user_name, secret)
    extra_params = {
        "user_name"     : "yanl",
        "client_version": "testing"
    }

    result = client.batch_get(gids=["ANDROID-ef550c2e5709eda4e51290dfeebda297",
                                    "ANDROID-e5315edf4bf94ed8b6ca6fda3ad2f7eb",
                                    "non-existent", ],
                              day="20200406",
                              features=["ft_usertags"],
                              extra_params=extra_params)

    index = 1
    print_result(index, result)


if __name__ == "__main__":
    demo_feature_client_troubleshooting(
        url="https://api-gateway.bi.getui.com/gt-feature-data-service/ordered/query-lite",
        user_name="ft-user",
        secret=generate_secret("user@2019"))
